import os
import shutil
from win32com import client
import pythoncom
import winreg
import time
import subprocess
import tkinter as tk
from tkinter import filedialog, messagebox
from colorama import init, Fore, Style
import sys
from tqdm import tqdm
import ctypes
import msvcrt
import struct

# 确保终端窗口显示
kernel32 = ctypes.WinDLL('kernel32')
kernel32.AllocConsole()
sys.stdout = open('CONOUT$', 'w', encoding='utf-8')
sys.stderr = open('CONOUT$', 'w', encoding='utf-8')
sys.stdin = open('CONIN$', 'r', encoding='utf-8')

# 初始化colorama
init()

def print_banner():
    """打印程序版权和说明信息"""
    banner = f"""
{Fore.CYAN}╔══════════════════════════════════════════════════════════════╗
║                    深圳中检联二维码替换工具                  ║
║                                                              ║
║                Copyright © 2025 by 草原企鹅 wechat:15242609  ║
╚══════════════════════════════════════════════════════════════╝{Style.RESET_ALL}

{Fore.YELLOW}使用说明：{Style.RESET_ALL}
1. 本程序用于批量替换Word文档中的二维码图片
2. 使用前请确保：
   - 源文档为Word格式（.docx）
   - 新二维码图片为PNG格式（.png）
   - 二维码图片的文件名与对应的Word文档名称相同
   例如：01.docx 对应 01.png

{Fore.YELLOW}注意事项：{Style.RESET_ALL}
1. 原始文档路径和输出路径不能为同一文件夹
2. 请确保文档未被其他程序占用
3. 建议处理前关闭所有Office程序
"""
    print(banner)

def select_folder(title):
    """打开文件夹选择对话框"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    folder = filedialog.askdirectory(title=title)
    return folder if folder else None

def confirm_paths(source_dir, qr_dir, output_dir):
    """确认路径信息"""
    print(f"\n{Fore.GREEN}已选择的路径：{Style.RESET_ALL}")
    print(f"源文档目录：{source_dir}")
    print(f"二维码目录：{qr_dir}")
    print(f"输出目录：{output_dir}")
    
    root = tk.Tk()
    root.withdraw()
    return messagebox.askyesno("确认", "确认使用这些路径吗？")

def kill_office_processes():
    """强制结束所有WPS和Word进程"""
    try:
        # 结束WPS进程
        subprocess.run(['taskkill', '/F', '/IM', 'wps.exe'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        subprocess.run(['taskkill', '/F', '/IM', 'wpsoffice.exe'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        # 结束Word进程
        subprocess.run(['taskkill', '/F', '/IM', 'winword.exe'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(1)  # 等待进程完全结束
    except:
        pass

def check_wps_installed():
    """检查是否安装了WPS"""
    try:
        # 检查WPS注册表
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Kingsoft\Office", 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY)
        winreg.CloseKey(key)
        return True
    except:
        return False

def check_word_installed():
    """检查是否安装了Microsoft Word"""
    try:
        # 检查Microsoft Office注册表
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Office\ClickToRun\Configuration", 0, winreg.KEY_READ | winreg.KEY_WOW64_64KEY)
        winreg.CloseKey(key)
        return True
    except:
        return False

def get_available_apps():
    """获取系统中可用的办公软件"""
    available_apps = []
    
    if check_wps_installed():
        available_apps.append("WPS")
    
    if check_word_installed():
        available_apps.append("Word")
    
    return available_apps

def check_word_version():
    """检查Word版本"""
    try:
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Office\14.0\Word\InstallRoot")
        return True
    except WindowsError:
        return False

def get_word_app(preferred_app=None):
    """获取指定的办公软件应用程序
    
    Args:
        preferred_app: 优先使用的软件，可以是 'WPS' 或 'Word'，
                      如果为 None，则让用户选择
    """
    # 先清理已有的进程
    kill_office_processes()
    
    available_apps = get_available_apps()
    
    if not available_apps:
        raise Exception("未安装 WPS 或 Microsoft Word")
    
    # 检查Word版本
    if check_word_version():
        print(
            "检测到您使用的是 Office 2010，可能无法正确处理某些新版本的Word文档。\n"
            "建议使用 WPS 进行处理。是否继续使用 Word？(y/n): "
        )
        choice = input().strip().lower()
        if choice != 'y':
            return get_wps_app()
    
    # 如果没有指定优先软件，且有多个可用软件，则让用户选择
    if preferred_app is None and len(available_apps) > 1:
        print(f"\n{Fore.CYAN}检测到多个可用的办公软件：{Style.RESET_ALL}")
        for i, app in enumerate(available_apps, 1):
            print(f"{i}. {app}")
        
        while True:
            try:
                choice = input(f"\n{Fore.YELLOW}请选择要使用的软件 (输入序号): {Style.RESET_ALL}")
                index = int(choice) - 1
                if 0 <= index < len(available_apps):
                    preferred_app = available_apps[index]
                    break
                else:
                    print(f"{Fore.RED}无效的选择，请重试{Style.RESET_ALL}")
            except ValueError:
                print(f"{Fore.RED}请输入有效的数字{Style.RESET_ALL}")
    # 如果指定了优先软件，但该软件不可用，则使用可用的软件
    elif preferred_app not in available_apps:
        preferred_app = available_apps[0]
    # 如果只有一个可用软件，直接使用
    elif preferred_app is None:
        preferred_app = available_apps[0]
    
    # 根据选择创建应用程序实例
    try:
        if preferred_app == "WPS":
            app = client.Dispatch("KWps.Application")
        else:
            app = client.Dispatch("Word.Application")
        
        # 设置应用程序属性
        app.Visible = False
        app.DisplayAlerts = False  # 禁用警告弹窗
        app.ScreenUpdating = False  # 禁用屏幕更新以提高性能
        
        return app, preferred_app
    except Exception as e:
        raise Exception(f"创建{preferred_app}实例失败: {str(e)}")

def test_open_document(app, doc_path):
    """测试是否能打开文档"""
    try:
        doc = app.Documents.Open(os.path.abspath(doc_path))
        doc.Close()
        return True
    except:
        return False

def check_doc_format(file_path):
    """检查Word文档的实际格式
    
    Returns:
        tuple: (actual_format, is_compound_doc)
        actual_format: 'doc' 或 'docx'
        is_compound_doc: 是否是复合文档
    """
    try:
        with open(file_path, 'rb') as f:
            header = f.read(8)
            
            # 检查是否是复合文档（DOC格式）
            if header.startswith(b'\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1'):
                return 'doc', True
            
            # 检查是否是ZIP文件（DOCX格式）
            if header.startswith(b'PK\x03\x04'):
                return 'docx', False
            
            return 'doc', True  # 如果无法确定，默认作为doc处理
    except:
        return 'doc', True  # 如果无法读取，默认作为doc处理

def open_document_safely(app, doc_path, app_type):
    """安全地打开文档，处理格式不匹配的情况"""
    actual_format, is_compound_doc = check_doc_format(doc_path)
    original_path = doc_path
    final_path = doc_path
    
    try:
        # 只有在使用Word且文件是doc格式但扩展名是docx时才进行转换
        if app_type == "Word" and actual_format == 'doc' and doc_path.lower().endswith('.docx'):
            try:
                # 先尝试直接打开docx
                doc = app.Documents.Open(os.path.abspath(doc_path))
                doc.Close()
                final_path = doc_path  # 可以打开docx，使用原路径
            except:
                # 如果打开失败，转换为doc格式
                final_path = doc_path[:-4] + 'doc'
                if original_path != final_path:
                    shutil.copy2(doc_path, final_path)
                    # 如果是在输出目录中，删除原始的docx文件
                    if os.path.dirname(original_path) == os.path.dirname(final_path):
                        os.remove(original_path)
        
        try:
            doc = app.Documents.Open(os.path.abspath(final_path))
            return doc, final_path
        except Exception as e:
            if app_type == "Word" and "WPS" in get_available_apps():
                raise Exception("Word无法打开文档，建议切换到WPS")
            raise e
    except:
        # 如果出错且生成了新文件，清理它
        if final_path != original_path and os.path.exists(final_path):
            try:
                os.remove(final_path)
            except:
                pass
        raise

def replace_qr_codes(source_dir, qr_dir, output_dir, preferred_app=None):
    """替换二维码的主函数"""
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有docx文件和二维码文件
    docx_files = [f for f in os.listdir(source_dir) if f.lower().endswith(('.docx', '.doc')) and not f.startswith('~$')]
    qr_files = {os.path.splitext(f)[0] for f in os.listdir(qr_dir) if f.lower().endswith('.png')}
    used_qr_files = set()
    
    if not docx_files:
        print(f"\n{Fore.RED}错误：在源目录中没有找到Word文档{Style.RESET_ALL}")
        return
    
    # 初始化 Word 应用程序
    word_app = None
    max_retries = 3
    
    try:
        with tqdm(total=len(docx_files), desc="处理进度", unit="文档") as pbar:
            for docx_file in docx_files:
                retries = 0
                success = False
                
                while retries < max_retries and not success:
                    try:
                        if word_app is None:
                            word_app, app_type = get_word_app(preferred_app)
                            if retries == 0:
                                print(f"\n{Fore.GREEN}使用 {app_type} 处理文档{Style.RESET_ALL}")
                        
                        # 检查对应的QR码图片是否存在
                        doc_name = os.path.splitext(docx_file)[0]
                        qr_image_path = os.path.join(qr_dir, doc_name + '.png')
                        if not os.path.exists(qr_image_path):
                            tqdm.write(f"{Fore.YELLOW}警告: 未找到对应的二维码图片: {doc_name}.png{Style.RESET_ALL}")
                            success = True
                            continue
                        
                        used_qr_files.add(doc_name)

                        # 获取完整路径
                        doc_path = os.path.join(source_dir, docx_file)
                        output_path = os.path.join(output_dir, docx_file)
                        
                        # 复制文档到输出目录
                        shutil.copy2(doc_path, output_path)
                        
                        # 打开文档
                        try:
                            doc, actual_output_path = open_document_safely(word_app, output_path, app_type)
                        except Exception as e:
                            if "建议切换到WPS" in str(e) and "WPS" in get_available_apps():
                                print(f"\n{Fore.YELLOW}Word无法打开文档，正在切换到WPS...{Style.RESET_ALL}")
                                word_app.Quit()
                                word_app = None
                                preferred_app = "WPS"
                                continue
                            else:
                                raise Exception(f"打开文档失败: {str(e)}")

                        try:
                            # 切换到页面布局视图
                            doc.ActiveWindow.View.Type = 3
                            
                            # 遍历所有浮动图片，找到第一页的二维码
                            found_qr = False
                            
                            # 找到Type=17的图片（二维码）
                            for i in range(1, doc.Shapes.Count + 1):
                                try:
                                    shape = doc.Shapes(i)
                                    shape.Select()
                                    
                                    # 检查是否在第一页
                                    if word_app.Selection.Information(3) == 1:
                                        if shape.Type == 17:  # 找到Type=17的图片
                                            # 使用ShapeRange来替换图片
                                            shape_range = doc.Shapes.Range([i])
                                            if app_type == "WPS":
                                                shape_range.Fill.UserPicture(os.path.abspath(qr_image_path))
                                            else:
                                                # 对于Word，直接替换图片
                                                shape.Fill.UserPicture(os.path.abspath(qr_image_path))
                                            
                                            found_qr = True
                                            break
                                except Exception as e:
                                    tqdm.write(f"{Fore.YELLOW}处理形状时出错: {str(e)}{Style.RESET_ALL}")
                                    continue
                            
                            if not found_qr:
                                tqdm.write(f"{Fore.YELLOW}警告: 文档 {docx_file} 中未找到二维码{Style.RESET_ALL}")
                            
                            # 保存文档
                            doc.Save()
                            success = True
                            
                        finally:
                            try:
                                doc.Close()
                            except:
                                pass
                        
                    except Exception as e:
                        if retries < max_retries - 1:
                            tqdm.write(f"{Fore.YELLOW}处理 {docx_file} 时出错，正在重试 ({retries + 1}/{max_retries}): {str(e)}{Style.RESET_ALL}")
                            if word_app:
                                try:
                                    word_app.Quit()
                                except:
                                    pass
                            kill_office_processes()
                            word_app = None
                            time.sleep(2)
                        else:
                            tqdm.write(f"{Fore.RED}处理 {docx_file} 失败: {str(e)}{Style.RESET_ALL}")
                        retries += 1
                
                pbar.update(1)
                
    finally:
        if word_app:
            try:
                word_app.DisplayAlerts = False
                word_app.ScreenUpdating = True
                word_app.Quit()
            except:
                pass
        kill_office_processes()
        
        # 检查未使用的二维码文件
        unused_qr_files = qr_files - used_qr_files
        if unused_qr_files:
            print(f"\n{Fore.YELLOW}以下二维码文件未被使用（可能是多余的）：{Style.RESET_ALL}")
            for qr_file in sorted(unused_qr_files):
                print(f"- {qr_file}.png")

def open_output_folder(folder_path):
    """打开输出文件夹"""
    try:
        os.startfile(os.path.abspath(folder_path))
    except:
        try:
            subprocess.run(['explorer', os.path.abspath(folder_path)])
        except:
            print(f"{Fore.YELLOW}无法自动打开输出文件夹，路径：{folder_path}{Style.RESET_ALL}")

def main():
    print_banner()
    
    # 选择路径
    print(f"\n{Fore.CYAN}请选择相关路径：{Style.RESET_ALL}")
    
    # 选择源文档目录
    print("\n请选择源文档所在目录...")
    source_dir = select_folder("选择源文档目录")
    if not source_dir:
        print(f"\n{Fore.RED}错误：未选择源文档目录，程序退出{Style.RESET_ALL}")
        return
    
    # 选择二维码图片目录
    print("\n请选择二维码图片所在目录...")
    qr_dir = select_folder("选择二维码图片目录")
    if not qr_dir:
        print(f"\n{Fore.RED}错误：未选择二维码图片目录，程序退出{Style.RESET_ALL}")
        return
    
    # 选择输出目录（不允许与源目录相同）
    while True:
        print("\n请选择处理后文档的保存目录...")
        output_dir = select_folder("选择输出目录")
        if not output_dir:
            print(f"\n{Fore.RED}错误：未选择输出目录，程序退出{Style.RESET_ALL}")
            return
        
        if os.path.normpath(output_dir) == os.path.normpath(source_dir):
            print(f"\n{Fore.RED}错误：输出目录不能与源文档目录相同，请重新选择{Style.RESET_ALL}")
            continue
        break
    
    # 确认路径
    if not confirm_paths(source_dir, qr_dir, output_dir):
        print(f"\n{Fore.YELLOW}已取消操作{Style.RESET_ALL}")
        return
    
    # 检测可用的办公软件
    available_apps = get_available_apps()
    if not available_apps:
        print(f"\n{Fore.RED}错误：未检测到安装的Office软件，请确保安装了WPS或Microsoft Office{Style.RESET_ALL}")
        print("\n按任意键退出...", end='', flush=True)
        msvcrt.getch()
        return
    
    # 如果有多个可用的办公软件，让用户选择
    preferred_app = None
    if len(available_apps) > 1:
        print(f"\n{Fore.CYAN}检测到多个可用的办公软件：{Style.RESET_ALL}")
        for i, app in enumerate(available_apps, 1):
            print(f"{i}. {app}")
        
        while True:
            try:
                choice = input(f"\n{Fore.YELLOW}请选择要使用的软件 (输入序号): {Style.RESET_ALL}")
                index = int(choice) - 1
                if 0 <= index < len(available_apps):
                    preferred_app = available_apps[index]
                    break
                else:
                    print(f"{Fore.RED}无效的选择，请重试{Style.RESET_ALL}")
            except ValueError:
                print(f"{Fore.RED}请输入有效的数字{Style.RESET_ALL}")
    else:
        preferred_app = available_apps[0]
        print(f"\n{Fore.GREEN}将使用 {preferred_app} 处理文档{Style.RESET_ALL}")
    
    print(f"\n{Fore.GREEN}开始处理...{Style.RESET_ALL}")
    
    # 初始化COM
    pythoncom.CoInitialize()
    try:
        replace_qr_codes(source_dir, qr_dir, output_dir, preferred_app)
    finally:
        pythoncom.CoUninitialize()
    
    print(f"\n{Fore.GREEN}处理完成！{Style.RESET_ALL}")
    print(f"\n处理后的文档已保存到：{output_dir}")
    
    # 打开输出文件夹并等待用户按键退出
    open_output_folder(output_dir)
    print("\n按任意键退出...", end='', flush=True)
    msvcrt.getch()

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(f"\n{Fore.RED}程序出错：{str(e)}{Style.RESET_ALL}")
        print("\n按任意键退出...", end='', flush=True)
        msvcrt.getch()